package org.cyk.solution.demo

import java.util.Date
import kotlin.math.min

class TokenBucket(
    private var rate: Int,  //令牌每秒生成数量(单位秒)
    private var capacity: Long,  //令牌桶容量
    private var tokens: Long = 0,  //当前令牌数量
    private var lastReqTime: Date = Date()  //处理上一个请求的时间
) {

    /**
     * 检查是否可以处理请求(获取令牌)
     */
    fun canConsumeToken(): Boolean {
        //计算上一个请求到这个到当前请求经过的时间
        val timeDiff = System.currentTimeMillis() - this.lastReqTime.time

        //根据时间差计算应该生成的令牌数
        val generatedTokensDouble = (timeDiff / 1000.0) * this.rate //timeDiff单位 -> 毫秒 | rate单位 -> 秒

        //对 generatedTokens 向上取整(令牌不能是小数)
        val generatedTokens = generatedTokensDouble.toLong()

        //更新令牌数量(不超过桶的容量)
        this.tokens = min(generatedTokens + this.tokens, this.capacity)

        //更新处理上一个请求的时间
        this.lastReqTime = Date(System.currentTimeMillis())

        //尝试消费一个令牌
        return this.tryConsumeToken()
    }

    /**
     * 尝试从桶中获取一个令牌
     */
    private fun tryConsumeToken(): Boolean =
        if(this.tokens > 0) {
            this.tokens -= 1
            true
        } else {
            false
        }

}

fun main() {
    //令牌桶创建的时候，默认上一个请求的到达的时间就是当前时间
    val tokenBucket = TokenBucket(
        rate = 4,
        capacity = 4
    )
    //模拟 900ms 时，同时来了四个请求
    Thread.sleep(900)
    val req1 = tokenBucket.canConsumeToken()
    val req2 = tokenBucket.canConsumeToken()
    val req3 = tokenBucket.canConsumeToken()
    val req4 = tokenBucket.canConsumeToken()
    //期望: 由于 1000ms 生成 4 个令牌，因此 900ms 的时候，只生成了 3 个令牌，前 3 个请求通过
    println("req1: $req1, req2: $req2, req3: $req3, req4: $req4")
}